#!/usr/bin/env python
import sys
import math
import numpy as np
import os, time, random
from multiprocessing import Pool

def process_images(name, task_cmd):
    print("=============== Checking env =============")
    print("[CBIG_CODE_DIR]: %s" % (os.getenv('CBIG_CODE_DIR')))
    print("[CBIG_MATLAB_DIR]: %s" % (os.getenv('CBIG_MATLAB_DIR')))
    print("[SUBJECTS_DIR]: %s" % (os.getenv('SUBJECTS_DIR')))
    print("[FSL_DIR]: %s" % (os.getenv('FSL_DIR')))
    print("[FSLTEMP_DIR]: %s" % (os.getenv('FSLTEMP_DIR')))
    print("[CBIG_AFNI_DIR]: %s" % (os.getenv('CBIG_AFNI_DIR')))
    print("[CBIG_ANTS_DIR]: %s" % (os.getenv('CBIG_ANTS_DIR')))
    print("================== Done ==================")

    # TODO: add env check failure instance

    print('Run task %s (%s)...' % (name, os.getpid()))
    start = time.time()
    print('[RUN CMD]: %s' % (task_cmd))
    os.system(task_cmd)
    end = time.time()
    print('Task %s runs %0.2f seconds.' % (name, (end - start)))

def get_task_infos(CBIG_preproc_dir, project_dir, src_task_file, processed_t1_dir, processed_rest_dir, cen_symbol):
    cmd_list = []

    list_dir = os.path.join(project_dir, 'Lists')
    ori_preproc_file = os.path.join(list_dir, 'preproc_config.txt')
    # generate config file
    preproc_file = generate_config_file(ori_preproc_file, cen_symbol)

    eval_csh_shell = os.path.join(CBIG_preproc_dir, "CBIG_preproc_fMRI_preprocess.csh")
    fp = open(src_task_file, "r")
    for line in fp.readlines():
        subj_name = line.replace("\n", "")
        # create fmri preporcess file
        fmri_conf_file = mk_preprocess_list_file(subj_name, project_dir)

        # -s paragram
        CBIG_preproc_cmd = eval_csh_shell + " -s " + subj_name
        CBIG_preproc_cmd = CBIG_preproc_cmd + " -output_d " + processed_rest_dir
        CBIG_preproc_cmd = CBIG_preproc_cmd + " -anat_s " + subj_name
        CBIG_preproc_cmd = CBIG_preproc_cmd + " -anat_d " + processed_t1_dir
        CBIG_preproc_cmd = CBIG_preproc_cmd + " -fmrinii " + fmri_conf_file
        CBIG_preproc_cmd = CBIG_preproc_cmd + " -config " + preproc_file
        cmd_list.append(CBIG_preproc_cmd)

    return cmd_list, len(cmd_list)

def generate_config_file(template_path, cen_symbol):
    print("########## Generate Config File ##########")

    line_list = []
    fp = open(template_path, 'r')
    for ori_line in fp.readlines():
        des_line = ori_line

        if 'skip' in ori_line:
            des_line = ori_line
            print("[Checking line]: Origin: `%s` --> Current: `%s` " % (ori_line, des_line))

        if 'fslslicetimer' in ori_line:
            full_file_path = ori_line.split(' ')[-1]
            file_path = full_file_path.split('/')[-1]
            file_type = file_path.split('.')[-1]

            dir_name = full_file_path.replace(file_path, "")
            file_name = file_path.replace('.' + file_type, "")

            tar_file_path = file_name + '_' + cen_symbol + '.' + file_type
            tar_full_file_path = dir_name + tar_file_path

            if not os.path.exists(tar_full_file_path.strip()):
                print("[ ## ERROR ## ]: File not exists: Could not find a file named `%s`" % (tar_full_file_path))
                print("[ ## NOTICE ## ]: Before you run this script, you need run `Scripts/generate_slice_order.py` first. ")
                sys.exit(0)

            ori_line_prefix = ori_line.replace(full_file_path, "")
            des_line = ori_line_prefix + tar_full_file_path

            print("[Checking line]: Origin: `%s` --> Current: `%s` " % (ori_line, des_line))

        if 'fslmcflirt_outliers' in ori_line:
            des_line = ori_line
            print("[Checking line]: Origin: `%s` --> Current: `%s` " % (ori_line, des_line))

        if 'bbregister' in ori_line:
            des_line = ori_line
            print("[Checking line]: Origin: `%s` --> Current: `%s` " % (ori_line, des_line))

        if 'regress' in ori_line:
            des_line = ori_line
            print("[Checking line]: Origin: `%s` --> Current: `%s` " % (ori_line, des_line))

        if 'censor' in ori_line:
            des_line = ori_line
            print("[Checking line]: Origin: `%s` --> Current: `%s` " % (ori_line, des_line))

        if 'bandpass' in ori_line:
            des_line = ori_line
            print("[Checking line]: Origin: `%s` --> Current: `%s` " % (ori_line, des_line))

        if 'QC_greyplot' in ori_line:
            des_line = ori_line
            print("[Checking line]: Origin: `%s` --> Current: `%s` " % (ori_line, des_line))

        if 'native2fsaverage' in ori_line:
            des_line = ori_line
            print("[Checking line]: Origin: `%s` --> Current: `%s` " % (ori_line, des_line))

        if 'FC_metrics' in ori_line:
            des_line = ori_line
            print("[Checking line]: Origin: `%s` --> Current: `%s` " % (ori_line, des_line))

        line_list.append(des_line)
    fp.close()

    file_type = template_path.split('.')[-1]
    des_name = template_path.replace("." + file_type, "")
    des_path = des_name + '_' + cen_symbol + '.' + file_type

    print("[Checking Done]: Finish Checking!")

    print('[Generate file]: Generate Config File...')
    fp = open(des_path, 'w')
    fp.writelines(line_list)
    fp.close()

    print('[Generate file]: Writing file in: `%s`' % (des_path))
    print('[Generate Done]: Finish Generate Config File!')
    print("########## Done ##########")

    return des_path

def mk_preprocess_list_file(subj_name, project_dir):
    # TODO: env variable
    CBIG_list_dir       = os.path.join(project_dir, 'Lists/CBIG_preprocess')
    origin_rest_dir     = os.path.join(project_dir, 'Origin/rest')
    subj_dir            = os.path.join(origin_rest_dir, subj_name)
    rest_dir_list      = os.listdir(subj_dir)

    rest_dir_list.sort()

    full_list_file_path = os.path.join(CBIG_list_dir, subj_name + ".txt")
    fp = open(full_list_file_path, "w")

    for file_index in range(len(rest_dir_list)):
        rest_dir_name = rest_dir_list[file_index]
        sess_mri_dir = os.path.join(subj_dir, rest_dir_name)

        file_name = os.listdir(sess_mri_dir)[0]
        full_fmri_file_path = os.path.join(sess_mri_dir, file_name)

        fmri_conf_str = "00" + str(file_index + 1) + " " + full_fmri_file_path
        fp.write(fmri_conf_str + '\n')
    fp.close()

    return full_list_file_path

def set_process_env():
    # set env
    os.environ["CBIG_CODE_DIR"] = "/home/li/Documents/Trainning/Share/Standalone_CBIG_fMRI_Preproc2016"
    os.environ["CBIG_MATLAB_DIR"] = "/usr/local/MATLAB/R2014b"
    os.environ["SUBJECTS_DIR"] = "/mri_projects/ASD_Analysis/DataProcessed/rest"
    os.environ["FSL_DIR"] = "/usr/lib/fsl/5.0"
    os.environ["FSLTEMP_DIR"] = "/usr/share/fsl"
    os.environ["CBIG_AFNI_DIR"] = "/usr/bin/afni"
    os.environ["CBIG_ANTS_DIR"] = "/usr/local/ants/bin"

if __name__=='__main__':
    print('Parent process %s.' % os.getpid())
    process_num = 11

    set_process_env()
    # TODO:
    symbol = "UCLA"
    target = "UCLA2"

    project_dir = '/mri_projects/ASD_Analysis'
    CBIG_preproc_dir = '/home/li/Documents/Trainning/Share/Standalone_CBIG_fMRI_Preproc2016/stable_projects/preprocessing/CBIG_fMRI_Preproc2016'
    # TODO:
    src_task_file = os.path.join(project_dir, 'Lists/recon_list_' + target + '.txt')
    processed_t1_dir = os.path.join(project_dir, 'DataProcessed/T1')
    processed_rest_dir = os.path.join(project_dir, 'DataProcessed/rest')

    src_lock_file = src_task_file + ".rest.lock"
    if os.path.exists(src_lock_file):
        print("[Notice]: File List `%s` Preporcessing Finished!" % (src_task_file))
    else:
        if not os.path.exists(processed_rest_dir):
            os.makedirs(processed_rest_dir)

        task_list, task_size = get_task_infos(CBIG_preproc_dir, project_dir, src_task_file, processed_t1_dir, processed_rest_dir, symbol)

        p = Pool(process_num)
        for i in range(task_size):
            p.apply_async(process_images, args=(i,task_list[i],))
        print('Waiting for all subprocesses done...')
        p.close()
        p.join()
        print('All subprocesses done.')

        print("########## Lock File ##########")
        print("[Lock rest]: %s" % (src_task_file))
        cp_cmd = "cp " + src_task_file + " " + src_lock_file
        os.system(cp_cmd)
        print("########## Done ##########")
